package com.example.mybatisplusproject.config;

import org.springframework.data.redis.cache.RedisCacheWriter;
import org.springframework.data.redis.connection.RedisConnection;
import org.springframework.data.redis.connection.RedisConnectionFactory;
import org.springframework.data.redis.core.Cursor;
import org.springframework.data.redis.core.ScanOptions;
import org.springframework.util.Assert;

import java.time.Duration;
import java.util.HashSet;
import java.util.Set;
import java.util.function.Function;

/**
 * 支持根据 key 设置缓存时间，不使用 keys 命令，nonLocking的 cache writer
 * @author zls
 * @date 2022/4/30 6:44 下午
 */
public class CustomRedisCacheWriter implements RedisCacheWriter {

    private final RedisCacheWriter redisCacheWriter;
    private final RedisConnectionFactory connectionFactory;

    public CustomRedisCacheWriter(RedisCacheWriter redisCacheWriter, RedisConnectionFactory factory) {
        this.redisCacheWriter = redisCacheWriter;
        this.connectionFactory = factory;
    }

    @Override
    public void put(String name, byte[] key, byte[] value, Duration ttl) {
        redisCacheWriter.put(name, key, value, getDuration(name, ttl));
    }

    /**
     * 根据缓存 key 的名称获取其有效时间
     * @param name
     * @param globalTtl
     * @return
     */
    private Duration getDuration(String name, Duration globalTtl) {
        String[] split = name.split("#");
        if (split.length == 1) {
            return globalTtl;
        }
        String s = split[split.length - 1];
        int minutes = Integer.parseInt(s);
        return Duration.ofMinutes(minutes);
    }

    @Override
    public byte[] get(String name, byte[] key) {
        return redisCacheWriter.get(name, key);
    }

    @Override
    public byte[] putIfAbsent(String name, byte[] key, byte[] value, Duration ttl) {
        return redisCacheWriter.putIfAbsent(name, key, value, getDuration(name, ttl));
    }

    @Override
    public void remove(String name, byte[] key) {
        redisCacheWriter.remove(name, key);
    }

    /**
     * 使用 scan 命令来替代 keys 命令
     * @param name
     * @param pattern
     */
    @Override
    public void clean(String name, byte[] pattern) {

        Assert.notNull(name, "Name must not be null!");
        Assert.notNull(pattern, "Pattern must not be null!");

        this.execute(connection -> {
            // 采用scan 来实现查找 key 的操作，不用 keys
            ScanOptions scanOptions = ScanOptions.scanOptions().match(new String(pattern)).count(1000).build();
            Cursor<byte[]> cursor = connection.scan(scanOptions);
            Set<byte[]> keys = new HashSet<>();
            cursor.forEachRemaining(keys::add);
            if (keys.size() > 0) {
                connection.del(keys.toArray(new byte[0][]));
            }
            return "OK";
        });
    }

    private <T> void execute(Function<RedisConnection, T> callback) {
        try (RedisConnection connection = connectionFactory.getConnection()) {
            callback.apply(connection);
        }
    }
}
